package process;

import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.KeyedProcessFunction;
import org.apache.flink.util.Collector;

public class KeyedProcessFunctionDemo {

    public static void main(String[] args) throws Exception{
        StreamExecutionEnvironment env = StreamExecutionEnvironment.createLocalEnvironmentWithWebUI(new Configuration());
        // 设置checkpoint
        env.enableCheckpointing(5000);

        // 读取数据 生成DataStream
        DataStreamSource<String> lines = env.socketTextStream("localhost", 8888);

        SingleOutputStreamOperator<Tuple2<String, Integer>> mapped = lines.map(new MapFunction<String, Tuple2<String, Integer>>() {
            @Override
            public Tuple2<String, Integer> map(String s) throws Exception {
                String[] fields = s.split(",");
                return Tuple2.of(fields[0], Integer.parseInt(fields[1]));
            }
        });

        mapped.keyBy(f -> f.f0)
                        .process(new MyKeyedProcessFunction())
                                .print();


        env.execute("");
    }

    private static class MyKeyedProcessFunction extends KeyedProcessFunction<String, Tuple2<String,Integer>,Tuple2<String,Integer>>{

        // 定义状态
        private transient ValueState<Integer> state;

        @Override
        public void open(Configuration parameters) throws Exception {
            // 创建状态描述器
            ValueStateDescriptor<Integer> stateDescriptor = new ValueStateDescriptor<>("value-state", Integer.class);
            // 获取状态
            state = getRuntimeContext().getState(stateDescriptor);
        }

        @Override
        public void processElement(Tuple2<String, Integer> value, Context context, Collector<Tuple2<String, Integer>> out) throws Exception {
            // 获取状态
            Integer historyValue = state.value();
            Integer currentValue = value.f1;
            Integer totalValue;
            if ( historyValue == null){
                historyValue = 0;
            }
            totalValue = currentValue + historyValue;
            state.update(totalValue);
            value.f1 = totalValue;
            out.collect(value);
        }
    }
}
